{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "057be575-8895-4335-8e72-d240c0d1c8e4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WebSocket server running on ws://172.24.197.154:8765\n",
      "New client connected: ('172.24.197.154', 51828)\n",
      "New client connected: ('172.24.197.154', 51832)\n",
      "[Client Message] 1\n",
      "[Client Message] 1\n",
      "[Client Message] 1\n",
      "New client connected: ('172.24.197.154', 51854)\n",
      "[Client Message] 1\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "[Server Input] >  2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New client connected: ('172.24.197.154', 51905)\n",
      "[Client Message] ok\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "[Server Input] >  你好\n"
     ]
    }
   ],
   "source": [
    "import socket\n",
    "import select\n",
    "import base64\n",
    "import hashlib\n",
    "import struct\n",
    "import threading\n",
    "\n",
    "# WebSocket 握手 GUID (RFC6455)\n",
    "WS_GUID = \"258EAFA5-E914-47DA-95CA-C5AB0DC85B11\"\n",
    "\n",
    "class WebSocketServer:\n",
    "    def __init__(self, host='172.24.197.154', port=8765):\n",
    "        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n",
    "        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n",
    "        self.server.bind((host, port))\n",
    "        self.server.listen(5)\n",
    "        self.clients = []\n",
    "        self.lock = threading.Lock()  # 新增客户端列表锁\n",
    "        \n",
    "    def generate_accept_key(self, key):\n",
    "        sha1 = hashlib.sha1((key + WS_GUID).encode()).digest()\n",
    "        return base64.b64encode(sha1).decode()\n",
    "    \n",
    "    def handle_handshake(self, client):\n",
    "        data = client.recv(4096).decode()\n",
    "        headers = dict([line.split(\": \", 1) for line in data.split(\"\\r\\n\")[1:-2]])\n",
    "        \n",
    "        if \"Sec-WebSocket-Key\" in headers:\n",
    "            accept_key = self.generate_accept_key(headers[\"Sec-WebSocket-Key\"])\n",
    "            response = (\n",
    "                \"HTTP/1.1 101 Switching Protocols\\r\\n\"\n",
    "                \"Upgrade: websocket\\r\\n\"\n",
    "                \"Connection: Upgrade\\r\\n\"\n",
    "                f\"Sec-WebSocket-Accept: {accept_key}\\r\\n\\r\\n\"\n",
    "            )\n",
    "            client.send(response.encode())\n",
    "            return True\n",
    "        return False\n",
    "    \n",
    "    def parse_frame(self, data):\n",
    "        # 解析 WebSocket 帧 (简化版，仅处理文本帧)\n",
    "        fin = data[0] & 0x80\n",
    "        opcode = data[0] & 0x0F\n",
    "        mask = data[1] & 0x80\n",
    "        payload_len = data[1] & 0x7F\n",
    "        \n",
    "        idx = 2\n",
    "        if payload_len == 126:\n",
    "            payload_len = struct.unpack(\"!H\", data[2:4])[0]\n",
    "            idx += 2\n",
    "        elif payload_len == 127:\n",
    "            payload_len = struct.unpack(\"!Q\", data[2:10])[0]\n",
    "            idx += 8\n",
    "            \n",
    "        if mask:\n",
    "            masking_key = data[idx:idx+4]\n",
    "            idx += 4\n",
    "            payload = data[idx:idx+payload_len]\n",
    "            decoded = bytearray([payload[i] ^ masking_key[i%4] for i in range(len(payload))])\n",
    "            return decoded.decode('utf-8')\n",
    "        return None\n",
    "    \n",
    "    def create_frame(self, message):\n",
    "        # 创建 WebSocket 帧 (服务端发送无需掩码)\n",
    "        header = bytearray()\n",
    "        payload = message.encode('utf-8')\n",
    "        header.append(0x81)  # FIN + Text Frame\n",
    "        if len(payload) <= 125:\n",
    "            header.append(len(payload))\n",
    "        elif len(payload) <= 65535:\n",
    "            header.append(126)\n",
    "            header.extend(struct.pack(\"!H\", len(payload)))\n",
    "        else:\n",
    "            header.append(127)\n",
    "            header.extend(struct.pack(\"!Q\", len(payload)))\n",
    "        return bytes(header) + payload\n",
    "    \n",
    "    def broadcast(self, sender, message):\n",
    "        frame = self.create_frame(message)\n",
    "        with self.lock:  # 加锁保护客户端列表\n",
    "            current_clients = list(self.clients)  # 创建副本避免遍历时修改\n",
    "            \n",
    "        for client in current_clients:\n",
    "            try:\n",
    "                if client != sender:  # 不发送给消息来源方\n",
    "                    client.send(frame)\n",
    "            except:\n",
    "                with self.lock:\n",
    "                    if client in self.clients:\n",
    "                        self.clients.remove(client)\n",
    "    \n",
    "    def console_input(self):\n",
    "        \"\"\"控制台输入线程函数\"\"\"\n",
    "        while True:\n",
    "            msg = input(\"[Server Input] > \")\n",
    "            if msg:\n",
    "                frame = self.create_frame(msg)\n",
    "                with self.lock:\n",
    "                    current_clients = list(self.clients)\n",
    "                for client in current_clients:\n",
    "                    try:\n",
    "                        client.send(frame)\n",
    "                    except:\n",
    "                        with self.lock:\n",
    "                            if client in self.clients:\n",
    "                                self.clients.remove(client)\n",
    "    \n",
    "    def run(self):\n",
    "        print(f\"WebSocket server running on ws://{self.server.getsockname()[0]}:{self.server.getsockname()[1]}\")\n",
    "        \n",
    "        # 启动控制台输入线程\n",
    "        input_thread = threading.Thread(target=self.console_input, daemon=True)\n",
    "        input_thread.start()\n",
    "        \n",
    "        while True:\n",
    "            rlist, _, _ = select.select([self.server] + self.clients, [], [])\n",
    "            \n",
    "            for sock in rlist:\n",
    "                if sock is self.server:  # 新连接\n",
    "                    client, addr = self.server.accept()\n",
    "                    if self.handle_handshake(client):\n",
    "                        with self.lock:\n",
    "                            self.clients.append(client)\n",
    "                        print(f\"New client connected: {addr}\")\n",
    "                else:  # 已有客户端消息\n",
    "                    try:\n",
    "                        data = sock.recv(4096)\n",
    "                        if data:\n",
    "                            message = self.parse_frame(data)\n",
    "                            if message:\n",
    "                                self.broadcast(sock, message)\n",
    "                                print(f\"[Client Message] {message}\")\n",
    "                        else:  # 连接关闭\n",
    "                            with self.lock:\n",
    "                                self.clients.remove(sock)\n",
    "                    except:\n",
    "                        with self.lock:\n",
    "                            if sock in self.clients:\n",
    "                                self.clients.remove(sock)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    server = WebSocketServer()\n",
    "    server.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "965feb94-3955-4a24-9277-0de22dca5606",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:base] *",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
